{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:08.832701Z",
     "start_time": "2025-07-10T10:15:08.828031Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from collections import Counter\n",
    "\n",
    "import jieba\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# 自定义简单数据集\n",
    "data = [\n",
    "    (\"你好，今天天气真好！\", \"Hello, the weather is nice today!\"),\n",
    "    (\"你吃饭了吗？\", \"Have you eaten yet?\"),\n",
    "    (\"深度学习很有趣。\", \"Deep learning is interesting.\"),\n",
    "    (\"我们一起学习吧。\", \"Let's study together.\"),\n",
    "    (\"这是一个测试例子。\", \"This is a test example.\")\n",
    "]\n",
    "\n",
    "\n",
    "# 中文分词函数\n",
    "def chinese_split(text):\n",
    "    return list(jieba.cut(text))  # 使用结巴分词\n",
    "\n",
    "\n",
    "# 英文分词函数\n",
    "def english_split(text):\n",
    "    return text.lower().split()\n",
    "\n",
    "\n",
    "# 处理原始数据\n",
    "chinese_sentences = [chinese_split(pair[0]) for pair in data]\n",
    "english_sentences = [english_split(pair[1]) for pair in data]\n",
    "\n",
    "chinese_sentences, english_sentences"
   ],
   "id": "ce053a7f01e08ecc",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([['你好', '，', '今天天气', '真', '好', '！'],\n",
       "  ['你', '吃饭', '了', '吗', '？'],\n",
       "  ['深度', '学习', '很', '有趣', '。'],\n",
       "  ['我们', '一起', '学习', '吧', '。'],\n",
       "  ['这是', '一个', '测试', '例子', '。']],\n",
       " [['hello,', 'the', 'weather', 'is', 'nice', 'today!'],\n",
       "  ['have', 'you', 'eaten', 'yet?'],\n",
       "  ['deep', 'learning', 'is', 'interesting.'],\n",
       "  [\"let's\", 'study', 'together.'],\n",
       "  ['this', 'is', 'a', 'test', 'example.']])"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 94
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:08.850728Z",
     "start_time": "2025-07-10T10:15:08.847388Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 处理特殊符号\n",
    "special_tokens = ['<pad>', '<bos>', '<eos>', '<unk>']\n",
    "\n",
    "\n",
    "# 构建词汇表\n",
    "def build_vocab(sentences):\n",
    "    counter = Counter()\n",
    "\n",
    "    for sentence in sentences:\n",
    "        for word in sentence:\n",
    "            counter[word] += 1\n",
    "\n",
    "    vocab = special_tokens.copy()\n",
    "\n",
    "    for word, count in counter.items():\n",
    "        if word not in special_tokens:\n",
    "            vocab.append(word)\n",
    "\n",
    "    word2idx = {word: idx for idx, word in enumerate(vocab)}\n",
    "    return vocab, word2idx\n",
    "\n",
    "\n",
    "# 构建中英文词汇表\n",
    "zh_vocab, zh_word2idx = build_vocab([sentence for sentence in chinese_sentences])\n",
    "\n",
    "en_vocab, en_word2idx = build_vocab([sentence for sentence in english_sentences])\n",
    "\n",
    "# print(zh_vocab)\n",
    "# print(en_vocab)\n",
    "print(zh_word2idx)\n",
    "print(en_word2idx)"
   ],
   "id": "ba851494e9bff76e",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'<pad>': 0, '<bos>': 1, '<eos>': 2, '<unk>': 3, '你好': 4, '，': 5, '今天天气': 6, '真': 7, '好': 8, '！': 9, '你': 10, '吃饭': 11, '了': 12, '吗': 13, '？': 14, '深度': 15, '学习': 16, '很': 17, '有趣': 18, '。': 19, '我们': 20, '一起': 21, '吧': 22, '这是': 23, '一个': 24, '测试': 25, '例子': 26}\n",
      "{'<pad>': 0, '<bos>': 1, '<eos>': 2, '<unk>': 3, 'hello,': 4, 'the': 5, 'weather': 6, 'is': 7, 'nice': 8, 'today!': 9, 'have': 10, 'you': 11, 'eaten': 12, 'yet?': 13, 'deep': 14, 'learning': 15, 'interesting.': 16, \"let's\": 17, 'study': 18, 'together.': 19, 'this': 20, 'a': 21, 'test': 22, 'example.': 23}\n"
     ]
    }
   ],
   "execution_count": 95
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:08.874445Z",
     "start_time": "2025-07-10T10:15:08.870034Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 参数设置\n",
    "ZH_VOCAB_SIZE = len(zh_vocab)\n",
    "EN_VOCAB_SIZE = len(en_vocab)\n",
    "HIDDEN_SIZE = 256\n",
    "BATCH_SIZE = 2\n",
    "LEARNING_RATE = 0.005\n",
    "\n",
    "\n",
    "def tokenize(words, word2idx):\n",
    "    # 如果某个词在字典中找不到，则用'<unk>'的索引代替\n",
    "    return [word2idx.get(word, word2idx['<unk>']) for word in words]\n",
    "\n",
    "\n",
    "processed_data_ch = []\n",
    "processed_data_en = []\n",
    "for ch, en in zip(chinese_sentences, english_sentences):\n",
    "    # ch，en分别是一个中文句子和对应的英文句子\n",
    "    # 在每个句子的前面加上'<bos>'，在每个句子的后面加上'<eos>'，这样大模型才能知道什么时候停止生成句子\n",
    "    ch_numerical = tokenize(ch, zh_word2idx)\n",
    "    en_numerical = [en_word2idx['<bos>']] + tokenize(en, en_word2idx) + [en_word2idx['<eos>']]\n",
    "    processed_data_ch.append(torch.LongTensor(ch_numerical))\n",
    "    processed_data_en.append(torch.LongTensor(en_numerical))\n",
    "\n",
    "# print(processed_data_ch)\n",
    "processed_data_en"
   ],
   "id": "bfc7c186f1217e14",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([1, 4, 5, 6, 7, 8, 9, 2]),\n",
       " tensor([ 1, 10, 11, 12, 13,  2]),\n",
       " tensor([ 1, 14, 15,  7, 16,  2]),\n",
       " tensor([ 1, 17, 18, 19,  2]),\n",
       " tensor([ 1, 20,  7, 21, 22, 23,  2])]"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 96
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:08.890720Z",
     "start_time": "2025-07-10T10:15:08.887987Z"
    }
   },
   "cell_type": "code",
   "source": [
    "processed_data_ch_pad = nn.utils.rnn.pad_sequence(processed_data_ch, batch_first=True,\n",
    "                                                  padding_value=zh_word2idx['<pad>'])\n",
    "processed_data_ch_pad"
   ],
   "id": "7e8909062d1e80f2",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 4,  5,  6,  7,  8,  9],\n",
       "        [10, 11, 12, 13, 14,  0],\n",
       "        [15, 16, 17, 18, 19,  0],\n",
       "        [20, 21, 16, 22, 19,  0],\n",
       "        [23, 24, 25, 26, 19,  0]])"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 97
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:08.929326Z",
     "start_time": "2025-07-10T10:15:08.926030Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 对processed_data进行数据填充，对齐长度\n",
    "processed_data_en_pad = nn.utils.rnn.pad_sequence(processed_data_en, batch_first=True,\n",
    "                                                  padding_value=en_word2idx['<pad>'])\n",
    "processed_data_en_pad"
   ],
   "id": "50c016488093de20",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1,  4,  5,  6,  7,  8,  9,  2],\n",
       "        [ 1, 10, 11, 12, 13,  2,  0,  0],\n",
       "        [ 1, 14, 15,  7, 16,  2,  0,  0],\n",
       "        [ 1, 17, 18, 19,  2,  0,  0,  0],\n",
       "        [ 1, 20,  7, 21, 22, 23,  2,  0]])"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 98
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:08.990122Z",
     "start_time": "2025-07-10T10:15:08.986894Z"
    }
   },
   "cell_type": "code",
   "source": [
    "\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "\n",
    "dataset = TensorDataset(processed_data_ch_pad, processed_data_en_pad)\n",
    "dataloader = DataLoader(dataset, batch_size=1, shuffle=False)\n",
    "\n",
    "for src, trg in dataloader:\n",
    "    print(src)\n",
    "    print(trg)\n",
    "    break"
   ],
   "id": "19524ded28178041",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[4, 5, 6, 7, 8, 9]])\n",
      "tensor([[1, 4, 5, 6, 7, 8, 9, 2]])\n"
     ]
    }
   ],
   "execution_count": 99
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:09.032411Z",
     "start_time": "2025-07-10T10:15:09.029402Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 位置编码\n",
    "# 大都督周瑜（我的微信: dadudu6789）\n",
    "\n",
    "import math\n",
    "\n",
    "\n",
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, max_seq_len=128):\n",
    "        super().__init__()\n",
    "        self.pe = torch.zeros(max_seq_len, d_model)\n",
    "        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        self.pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        self.pe[:, 1::2] = torch.cos(position * div_term)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.pe[:x.size(1), :]\n",
    "        return x"
   ],
   "id": "3c9ad1f6b832098a",
   "outputs": [],
   "execution_count": 100
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:09.067193Z",
     "start_time": "2025-07-10T10:15:09.063646Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class ZhouyuModel(nn.Module):\n",
    "    def __init__(self, d_model, dim_feedforward, nhead, num_encoder_layers, num_decoder_layers):\n",
    "        super().__init__()\n",
    "        self.encoder_embedding = nn.Embedding(len(zh_vocab), d_model)\n",
    "        self.decoder_embedding = nn.Embedding(len(en_vocab), d_model)\n",
    "        self.pos_encoder = PositionalEncoding(d_model)\n",
    "        self.transformer = nn.Transformer(d_model=d_model, dim_feedforward=dim_feedforward, nhead=nhead,\n",
    "                                          num_encoder_layers=num_encoder_layers,\n",
    "                                          num_decoder_layers=num_decoder_layers, batch_first=True, dropout=0)\n",
    "        self.fc = nn.Linear(d_model, len(en_vocab))\n",
    "\n",
    "    def forward(self, zh_inputs, en_inputs):\n",
    "\n",
    "        batch_size, en_seq_len = en_inputs.shape\n",
    "        # mask = torch.tril(torch.ones(en_seq_len, en_seq_len))\n",
    "        mask = nn.Transformer.generate_square_subsequent_mask(en_seq_len)\n",
    "\n",
    "        # 词嵌入和位置编码\n",
    "        encoder_input = self.pos_encoder(self.encoder_embedding(zh_inputs))\n",
    "        decoder_input = self.pos_encoder(self.decoder_embedding(en_inputs))\n",
    "\n",
    "        output = self.transformer(\n",
    "            src=encoder_input, tgt=decoder_input,\n",
    "            tgt_mask=mask\n",
    "        )\n",
    "\n",
    "        return self.fc(output)"
   ],
   "id": "b9f10077b71810ed",
   "outputs": [],
   "execution_count": 101
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:09.102005Z",
     "start_time": "2025-07-10T10:15:09.085291Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model = ZhouyuModel(d_model=128, dim_feedforward=2048, nhead=8, num_encoder_layers=2, num_decoder_layers=2)\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=en_word2idx['<pad>'])"
   ],
   "id": "8d5734d16acad79e",
   "outputs": [],
   "execution_count": 102
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:09.120004Z",
     "start_time": "2025-07-10T10:15:09.117354Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 打印一下model的参数个数\n",
    "print(sum(p.numel() for p in model.parameters()))"
   ],
   "id": "b44ae3b38a67757a",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2514840\n"
     ]
    }
   ],
   "execution_count": 103
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:10.424980Z",
     "start_time": "2025-07-10T10:15:09.137285Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for epoch in range(20):\n",
    "    # input： tensor([[4, 5, 6, 7, 8, 9]])\n",
    "    # target：tensor([[1, 4, 5, 6, 7, 8, 9, 2]])\n",
    "    for input, target in dataloader:\n",
    "        # 准备解码器输入输出，其实这一步可以在数据处理时做掉\n",
    "        decoder_input = target[:, :-1]  # 移除最后一个token  tensor([[1, 4, 5, 6, 7, 8, 9]])\n",
    "        decoder_target = target[:, 1:]  # 移除第一个token    tensor([[4, 5, 6, 7, 8, 9, 2]])\n",
    "\n",
    "        decoder_outputs = model(input, decoder_input)\n",
    "\n",
    "        # 计算损失\n",
    "        loss = criterion(\n",
    "            # decoder_output本来是(batch_size, seq_len, vocab_size)，变成(batch_size * seq_len, vocab_size)\n",
    "            # decoder_target本来是(batch_size, seq_len)，变成(batch_size * seq_len)\n",
    "            decoder_outputs.view(-1, decoder_outputs.size(-1)),\n",
    "            decoder_target.view(-1)\n",
    "        )\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        print(f'Epoch {epoch + 1}, Loss: {loss:.4f}')"
   ],
   "id": "2573befd839be55c",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 3.3264\n",
      "Epoch 1, Loss: 3.7421\n",
      "Epoch 1, Loss: 3.7243\n",
      "Epoch 1, Loss: 4.0434\n",
      "Epoch 1, Loss: 3.9274\n",
      "Epoch 2, Loss: 2.8441\n",
      "Epoch 2, Loss: 2.6646\n",
      "Epoch 2, Loss: 2.3883\n",
      "Epoch 2, Loss: 2.7236\n",
      "Epoch 2, Loss: 3.2953\n",
      "Epoch 3, Loss: 2.5373\n",
      "Epoch 3, Loss: 2.4220\n",
      "Epoch 3, Loss: 1.8835\n",
      "Epoch 3, Loss: 1.9247\n",
      "Epoch 3, Loss: 2.4776\n",
      "Epoch 4, Loss: 1.0088\n",
      "Epoch 4, Loss: 1.3849\n",
      "Epoch 4, Loss: 1.2230\n",
      "Epoch 4, Loss: 1.0992\n",
      "Epoch 4, Loss: 1.6393\n",
      "Epoch 5, Loss: 0.7535\n",
      "Epoch 5, Loss: 0.9344\n",
      "Epoch 5, Loss: 0.7792\n",
      "Epoch 5, Loss: 0.6252\n",
      "Epoch 5, Loss: 1.4751\n",
      "Epoch 6, Loss: 0.5453\n",
      "Epoch 6, Loss: 0.4671\n",
      "Epoch 6, Loss: 0.6886\n",
      "Epoch 6, Loss: 0.6248\n",
      "Epoch 6, Loss: 0.5865\n",
      "Epoch 7, Loss: 0.3908\n",
      "Epoch 7, Loss: 0.4004\n",
      "Epoch 7, Loss: 0.5349\n",
      "Epoch 7, Loss: 0.4084\n",
      "Epoch 7, Loss: 0.4034\n",
      "Epoch 8, Loss: 0.3308\n",
      "Epoch 8, Loss: 0.2888\n",
      "Epoch 8, Loss: 0.3510\n",
      "Epoch 8, Loss: 0.2535\n",
      "Epoch 8, Loss: 0.3579\n",
      "Epoch 9, Loss: 0.1949\n",
      "Epoch 9, Loss: 0.2005\n",
      "Epoch 9, Loss: 0.2476\n",
      "Epoch 9, Loss: 0.1317\n",
      "Epoch 9, Loss: 0.1958\n",
      "Epoch 10, Loss: 0.1357\n",
      "Epoch 10, Loss: 0.1157\n",
      "Epoch 10, Loss: 0.1599\n",
      "Epoch 10, Loss: 0.0894\n",
      "Epoch 10, Loss: 0.0979\n",
      "Epoch 11, Loss: 0.1076\n",
      "Epoch 11, Loss: 0.0825\n",
      "Epoch 11, Loss: 0.0874\n",
      "Epoch 11, Loss: 0.0654\n",
      "Epoch 11, Loss: 0.0662\n",
      "Epoch 12, Loss: 0.0707\n",
      "Epoch 12, Loss: 0.0580\n",
      "Epoch 12, Loss: 0.0653\n",
      "Epoch 12, Loss: 0.0452\n",
      "Epoch 12, Loss: 0.0508\n",
      "Epoch 13, Loss: 0.0490\n",
      "Epoch 13, Loss: 0.0426\n",
      "Epoch 13, Loss: 0.0449\n",
      "Epoch 13, Loss: 0.0357\n",
      "Epoch 13, Loss: 0.0402\n",
      "Epoch 14, Loss: 0.0392\n",
      "Epoch 14, Loss: 0.0338\n",
      "Epoch 14, Loss: 0.0359\n",
      "Epoch 14, Loss: 0.0299\n",
      "Epoch 14, Loss: 0.0342\n",
      "Epoch 15, Loss: 0.0327\n",
      "Epoch 15, Loss: 0.0283\n",
      "Epoch 15, Loss: 0.0305\n",
      "Epoch 15, Loss: 0.0256\n",
      "Epoch 15, Loss: 0.0299\n",
      "Epoch 16, Loss: 0.0280\n",
      "Epoch 16, Loss: 0.0248\n",
      "Epoch 16, Loss: 0.0272\n",
      "Epoch 16, Loss: 0.0230\n",
      "Epoch 16, Loss: 0.0261\n",
      "Epoch 17, Loss: 0.0251\n",
      "Epoch 17, Loss: 0.0223\n",
      "Epoch 17, Loss: 0.0250\n",
      "Epoch 17, Loss: 0.0210\n",
      "Epoch 17, Loss: 0.0235\n",
      "Epoch 18, Loss: 0.0232\n",
      "Epoch 18, Loss: 0.0207\n",
      "Epoch 18, Loss: 0.0231\n",
      "Epoch 18, Loss: 0.0196\n",
      "Epoch 18, Loss: 0.0217\n",
      "Epoch 19, Loss: 0.0217\n",
      "Epoch 19, Loss: 0.0193\n",
      "Epoch 19, Loss: 0.0215\n",
      "Epoch 19, Loss: 0.0184\n",
      "Epoch 19, Loss: 0.0203\n",
      "Epoch 20, Loss: 0.0204\n",
      "Epoch 20, Loss: 0.0182\n",
      "Epoch 20, Loss: 0.0203\n",
      "Epoch 20, Loss: 0.0175\n",
      "Epoch 20, Loss: 0.0190\n"
     ]
    }
   ],
   "execution_count": 104
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-10T10:15:13.110313Z",
     "start_time": "2025-07-10T10:15:13.098497Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 翻译函数\n",
    "def translate(sentence, model):\n",
    "    zh_tokens = torch.LongTensor(tokenize(chinese_split(sentence), zh_word2idx))\n",
    "    encoder_input = model.pos_encoder(model.encoder_embedding(zh_tokens).unsqueeze(0))\n",
    "    encoder_outputs = model.transformer.encoder(encoder_input)\n",
    "\n",
    "    decoder_inputs = [en_word2idx['<bos>']]\n",
    "\n",
    "    for _ in range(50):\n",
    "        with torch.no_grad():\n",
    "            decoder_input = model.pos_encoder(model.decoder_embedding(torch.LongTensor(decoder_inputs)).unsqueeze(0))\n",
    "            decoder_output = model.transformer.decoder(tgt=decoder_input, memory=encoder_outputs, tgt_mask=None)\n",
    "            output = model.fc(decoder_output)\n",
    "            pred_token = output[:, -1, :].argmax().item()\n",
    "            decoder_inputs.append(pred_token)\n",
    "            if pred_token == en_word2idx['<eos>']:\n",
    "                break\n",
    "\n",
    "    return ' '.join([en_vocab[idx] for idx in decoder_inputs[1:-1]])\n",
    "\n",
    "\n",
    "# 测试翻译\n",
    "test_sentence = \"我们一起学习\"\n",
    "print(translate(test_sentence, model))"
   ],
   "id": "69fa2a11cfd17fcf",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "let's study together.\n"
     ]
    }
   ],
   "execution_count": 110
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
